import os

from PIL import Image, ImageOps
import numpy as np
from tqdm import tqdm

# 定义标签到RGB颜色的映射表
# 请根据你的实际标签和颜色进行调整

label_to_color = {
    0: (0, 0, 0),  # background
    1: (8, 248, 0),  # grass
    2: (208, 244, 247),  # road
    3: (193, 95, 96),  # farmland
    4: (89, 164, 255),  # river
    5: (42, 14, 255),  # lake
    6: (148, 82, 255),  # container
    7: (255, 65, 220),  # building
    8: (8, 127, 0),  # forest
    9: (103, 27, 4),  # bare-soil
    10: (118, 128, 164),  # paved-surface
    11: (255, 108, 16),  # railway
    12: (140, 41, 133),  # greenhouse
    13: (255, 88, 98),  # stadium
    # 添加更多标签和颜色
}


def label_to_rgb(mask, colormap):
    h, w = mask.shape
    color_image = np.zeros((h, w, 3), dtype=np.uint8)

    for label, color in colormap.items():
        color_image[mask == label] = color

    return Image.fromarray(color_image)


# 加载图片和mask图像
image_path = 'path_to_your_image.jpg'

ouput_mask_folder = r"G:\SGG\segmmetation\data_enforce\imagePNG"
output_color_folder = r"G:\SGG\segmmetation\data_enforce\imagecolor"
os.makedirs(output_color_folder,exist_ok=True)
for mask_file in tqdm(os.listdir(ouput_mask_folder), desc="Process img files", unit='files'):
    mask_path = os.path.join(ouput_mask_folder, mask_file)
    mask = Image.open(mask_path).convert('L')  # 确保mask是灰度图

    assert mask.size == (512, 512)

    # 将mask转换为numpy数组
    mask_array = np.array(mask, dtype=np.int8)

    # 将mask标签映射到RGB颜色
    color_mask = label_to_rgb(mask_array, label_to_color)
    save_color_path = os.path.join(output_color_folder,mask_file)
    # 保存彩色mask图像
    color_mask.save(save_color_path)

print("mask图像转换为RGB颜色并保存为color.png。")
